import matplotlib.pyplot as plt
import numpy as np
import os
from datetime import datetime


# 从文件中读取损失数据
def read_loss_data(file_path):
    losses = []
    with open(file_path, "r") as file:
        for line in file:
            if "loss in epoch" in line:
                # 提取损失值
                loss = float(line.strip().split(": ")[-1])
                losses.append(loss)
    return losses


# 创建输出目录
output_dir = "ml-1m_default"
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 使用当前时间创建文件名
current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
output_file = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    f"performance_comparison_{current_time}.png",
)


# 读取数据
# TODO 这里需要改地址
infini_transformer_losses = read_loss_data(
    "/Users/zzq/Developer/python/SAM/reproduction/infini-transformer/myCode/SASRec-in-InfiniTransformer/ml-1m_default/infiniTransformer_20250102_205318.txt"
)

routing_transformer_losses = read_loss_data(
    "/Users/zzq/Developer/python/SAM/sasrec_ex/myCode/ml-1m_output/routingTransformer_20250102_144322.txt"
)

transformer_losses = read_loss_data(
    "/Users/zzq/Developer/python/SAM/sasrec_ex/myCode/ml-1m_output/transformer-output.txt"
)

# 绘制运行时间和准确性的对比图
plt.figure(figsize=(18, 12))

# 绘制损失对比图
plt.subplot(3, 1, 1)
plt.plot(
    infini_transformer_losses,
    label="Infini Transformer",
    linewidth=2,
    linestyle="-",
    marker="o",
    markersize=4,
    markevery=5,
)
plt.plot(
    routing_transformer_losses,
    label="Routing Transformer",
    linewidth=2,
    linestyle="--",
    marker="s",
    markersize=4,
    markevery=5,
)
plt.plot(
    transformer_losses,
    label="Simple Transformer",
    linewidth=2,
    linestyle=":",
    marker="^",
    markersize=4,
    markevery=5,
)
plt.xlabel("Iteration")
plt.ylabel("Loss")
plt.title("Loss Comparison")
plt.grid(True, linestyle="--", alpha=0.7)
plt.legend(loc="upper right")

# TODO 这里需要改参数 绘制运行时间对比和绘制准确性对比 前面的是Infini
# epoch:100, time: 1361.599287(s), valid (NDCG@10: 0.5878, HR@10: 0.8242), test (NDCG@10: 0.5633, HR@10: 0.8018)
# epoch:100, time: 1407.427627(s), valid (NDCG@10: 0.5774, HR@10: 0.8131), test (NDCG@10: 0.5536, HR@10: 0.7871)

# infini_transformer_losses
# epoch:100, time: 1142.217348(s), valid (NDCG@10: 0.6089, HR@10: 0.8406), test (NDCG@10: 0.5835, HR@10: 0.8157)
valid_time_values = [1142.217348, 1407.427627, 1796.537418]

valid_ndcg_values = [0.6089, 0.5774, 0.6117]
valid_hr_values = [0.8406, 0.8131, 0.8455]
test_ndcg_values = [0.5835, 0.5536, 0.5838]
test_hr_values = [0.8157, 0.7871, 0.8157]

plt.subplot(3, 1, 2)
plt.bar(
    ["Infini Transformer", "Routing Transformer", "Simple Transformer"],
    valid_time_values,
    # color=["lightblue", "salmon", "lightgreen"]
)
plt.ylabel("Time (seconds)")
plt.title("Running Time Comparison")


plt.subplot(3, 1, 3)
bar_width = 0.2
index = np.arange(len(valid_ndcg_values))

plt.bar(
    index,
    valid_ndcg_values,
    bar_width,
    label="Valid NDCG@10",
    color="lightgreen",
    alpha=0.7,
)
plt.bar(
    index + bar_width,
    valid_hr_values,
    bar_width,
    label="Valid HR@10",
    color="orange",
    alpha=0.7,
)
plt.bar(
    index + bar_width * 2,
    test_ndcg_values,
    bar_width,
    label="Test NDCG@10",
    color="darkgreen",
    alpha=0.7,
)
plt.bar(
    index + bar_width * 3,
    test_hr_values,
    bar_width,
    label="Test HR@10",
    color="red",
    alpha=0.7,
)
plt.xlabel("Model")
plt.ylabel("Score")
plt.title("Accuracy Metrics Comparison")
plt.xticks(
    index + bar_width * 1.5,
    ["Infini Transformer", "Routing Transformer", "Transformer"],
)
plt.legend()

plt.tight_layout()
plt.savefig(output_file)
plt.show()
